from concurrent.futures import ThreadPoolExecutor
from embedding_onnx_provider import ONNXExecutionProvider
from typing import List
from log_config import logger

class ThreadPoolEngine:

    def __init__(self,workers) -> None:
        self.running = False
        self.workers = workers

    def start(self):
        if self.running:
            logger.warning("该引擎已经在运行中")
            return
        
        self.thread_pool = ThreadPoolExecutor(max_workers=self.workers)
        self.running = True
        logger.info("Engine started.")
    
    def stop(self):
        self._check_running()
        self.running = False
        if self.thread_pool is not None:
            self.thread_pool.shutdown(wait=True)
        logger.info("Engine stoped.")

    def _check_running(self):
        if not self.running:
            raise ValueError("The engine is not running. ""You must start the engine before using it.")
        
    def _submit(self,fn, **kwargs):
        self._check_running()
        if self.thread_pool is None:
            raise ValueError("ThreadPoolExecutor is not initialized.")
        return self.thread_pool.submit(fn,**kwargs)


class EmbeddingEngine(ThreadPoolEngine):

    def __init__(self, embed_path, workers):
        super().__init__(workers)
        self.executor = ONNXExecutionProvider(embed_path)

    def execute(self, sentences: List[str]):
        task = self._submit(self.executor.encode,sentences=sentences)
        return task.result()